import os
import torch
import numpy as np
from torch.utils.data import DataLoader

from ENF_deepLearning.utils.cnn.Net_Dataset import Net_Dataset
from ENF_deepLearning.utils.helper.Data_Helper import Data_Helper
from ENF_deepLearning.utils.helper.File_Helper import File_Helper


class RNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = torch.nn.LSTM(
            input_size=35000,
            hidden_size=128,
            num_layers=1,
            batch_first=True
        )
        self.out = torch.nn.Linear(in_features=128, out_features=2)

    def forward(self, x):
        # 以下关于shape的注释只针对单向
        # output: [batch_size, time_step, hidden_size]
        # h_n: [num_layers,batch_size, hidden_size] # 虽然LSTM的batch_first为True,但是h_n/c_n的第一维还是num_layers
        # c_n: 同h_n
        output, (h_n, c_n) = self.rnn(x)
        # print(output.size())
        # output_in_last_timestep=output[:,-1,:] # 也是可以的
        output_in_last_timestep = h_n[-1, :, :]
        # print(output_in_last_timestep.equal(output[:,-1,:])) #ture
        x = self.out(output_in_last_timestep)
        return x


if __name__ == "__main__":

    file_helper = File_Helper()  # 文件预处理有关的操作助手
    data_helper = Data_Helper()  # 数据相关的操作助手
    data_path = '../utils/data'
    list = os.listdir(data_path)

    Y = []
    X = []
    minn = 0
    size = 35000
    PAD = size + 1
    for item in list:
        path = data_path + '/' + item
        data = np.loadtxt(path, dtype=np.float, delimiter=',')
        minn = max(minn, data.shape[0])
        data = data.tolist()
        if len(data) < PAD:
            for i in range(PAD - len(data)):
                data.append(0)
        # plt.plot(data)
        # plt.show()
        X.append(np.array(data[1:]).astype(np.float))
        Y.append(np.array(data[0]).astype(np.int))

    X = np.array(X)
    Y = np.array(Y)

    train_size = int(X.shape[0] * 0.7)
    test_size = len(X) - train_size
    train_X = X[:train_size]
    train_Y = Y[:train_size]
    test_X = X[train_size:]
    test_Y = Y[train_size:]

    train_X = train_X.reshape(-1, 1, size)  # 输入层维度必须为三维

    ss = Net_Dataset(train_X, train_Y)
    dataloader = DataLoader(dataset=ss, batch_size=32, shuffle=True, num_workers=2)


    test_data = Net_Dataset(test_X, test_Y)
    test_dataloader = DataLoader(dataset=test_data, batch_size=1000, shuffle=False, num_workers=2)

    testdata_iter = iter(test_dataloader)
    test_x, test_y = testdata_iter.next()
    test_x = test_x.view(-1, 1, size)
    test_x = test_x.to('cuda')
    test_y = test_y.to('cuda')
    # 2. 网络搭建
    net = RNN().to('cuda')
    # 3. 训练
    # 3. 网络的训练（和之前CNN训练的代码基本一样）
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
    loss_F = torch.nn.CrossEntropyLoss()
    for epoch in range(1000):  # 数据集只迭代一次
        for step, input_data in enumerate(dataloader):
            x, y = input_data
            x = x.to('cuda')
            y = y.to('cuda')
            pred = net(x.view(-1, 1, size))

            loss = loss_F(pred, y)  # 计算loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        if (epoch + 1) % 2 == 0:  # 每 100 次输出结果
            with torch.no_grad():
                test_pred = net(test_x)
                prob = torch.nn.functional.softmax(test_pred, dim=1)
                pred_cls = torch.argmax(prob, dim=1)
                acc = (pred_cls == test_y).sum().float().cpu().numpy() / pred_cls.size()[0]
                print('epoch:', epoch + 1, 'loss:', loss.item(), 'test acc:', acc)
